#ifndef WORKTHREADPOOL_H_
#define WORKTHREADPOOL_H_

#include <functional>
#include <memory>
#include <atomic>
#include <vector>
#include <unordered_map>
#include <mutex>
#include <thread>
#include "TaskQueue.h"
#include "ThreadGroup.h"
#include "TaskExecutor.h"

namespace toolkit {

/**
 * @brief 工作线程类，每个线程拥有自己的任务队列
 */
class WorkThread : public TaskExecutor {
public:
    using Ptr = std::shared_ptr<WorkThread>;

    WorkThread();
    ~WorkThread() override;

    void start();

    void stop();
    
    /**
     * @brief 异步执行任务
     * @param task 任务函数
     * @param may_sync 是否允许同步执行
     * @return 任务取消对象
     */
    TaskCancelable::Ptr async(Task task, bool may_sync = true) override;
    
    /**
     * @brief 异步执行任务，优先执行
     * @param task 任务函数
     * @param may_sync 是否允许同步执行
     * @return 任务取消对象
     */
    TaskCancelable::Ptr async_first(Task task, bool may_sync = true) override;
    
    /**
     * @brief 获取任务执行器负载
     * @return 负载值
     */
    int getLoad() override;
    
    /**
     * @brief 获取线程ID
     * @return 线程ID
     */
    std::thread::id getThreadId() const;
    
    /**
     * @brief 判断当前线程是否是工作线程
     * @return 是否是工作线程
     */
    bool isCurrentThread() const;
    
private:
    /**
     * @brief 线程函数
     */
    void run();
    
private:
    std::atomic<bool> _exit_flag{false};
    std::thread _thread;
    std::thread::id _thread_id;
    ThreadLoadCounter _thread_load;
    TaskQueue<Task> _task_queue;
};

/**
 * @brief 工作线程池类，每个线程有自己的任务队列，避免竞争
 */
class WorkThreadPool : public TaskExecutor {
public:
    using Ptr = std::shared_ptr<WorkThreadPool>;
    
    /**
     * @brief 构造函数
     * @param thread_num 线程数量
     */
    WorkThreadPool(size_t thread_num = 1);
    
    ~WorkThreadPool() override;
    
    /**
     * @brief 获取线程数量
     */
    size_t size() const;
    
    /**
     * @brief 异步执行任务
     * @param task 任务函数
     * @param may_sync 是否允许同步执行
     * @return 任务取消对象
     */
    TaskCancelable::Ptr async(Task task, bool may_sync = true) override;
    
    /**
     * @brief 异步执行任务，优先执行
     * @param task 任务函数
     * @param may_sync 是否允许同步执行
     * @return 任务取消对象
     */
    TaskCancelable::Ptr async_first(Task task, bool may_sync = true) override;
    
    /**
     * @brief 获取任务执行器负载
     * @return 负载值
     */
    int getLoad() override;
    
    /**
     * @brief 获取指定索引的工作线程
     * @param index 线程索引
     * @return 工作线程
     */
    WorkThread::Ptr getWorkThread(size_t index);
    
    /**
     * @brief 获取当前线程对应的工作线程
     * @return 工作线程，如果当前线程不是工作线程则返回nullptr
     */
    WorkThread::Ptr getCurrentWorkThread();
    
    /**
     * @brief 获取负载最小的工作线程
     * @return 负载最小的工作线程
     */
    WorkThread::Ptr getMinLoadWorkThread();
    
private:
    size_t _thread_num;
    std::vector<WorkThread::Ptr> _threads;
    std::mutex _mutex;
    std::unordered_map<std::thread::id, WorkThread::Ptr> _thread_map;
};

} /* namespace toolkit */

#endif /* WORKTHREADPOOL_H_ */